import pickle
import tarfile

import numpy as np
from matplotlib import pyplot as plt

from core.data.dataset import Dataset
from implement.transforms.astype import ToFloat
from implement.transforms.compose import Compose
from implement.transforms.normalize import Normalize
from utils.file_opera import get_file, save_cache_npz, load_cache_npz


class CIFAR10(Dataset):
    """
    CIFAR-10 数据集类。
    """

    def __init__(self, train=True,
                 transform=Compose([ToFloat(), Normalize(mean=0.5, std=0.5)]),
                 target_transform=None):
        """
        初始化 CIFAR-10 数据集。

        Args:
            train (bool, optional): 是否为训练集。默认为 True。
            transform (callable, optional): 数据变换函数。默认为标准化。
            target_transform (callable, optional): 标签变换函数。默认为 None。
        """
        super().__init__(train, transform, target_transform)

    def prepare(self):
        """
        准备 CIFAR-10 数据集。
        """
        url = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'
        self.data, self.label = load_cache_npz(url, self.train)
        if self.data is not None:
            return
        filepath = get_file(url)
        if self.train:
            self.data = np.empty((50000, 3 * 32 * 32))
            self.label = np.empty((50000), dtype=int)
            for i in range(5):
                self.data[i * 10000:(i + 1) * 10000] = self._load_data(
                    filepath, 'train', i + 1)
                self.label[i * 10000:(i + 1) * 10000] = self._load_label(
                    filepath, 'train', i + 1)
        else:
            self.data = self._load_data(filepath, 'test', 5)
            self.label = self._load_label(filepath, 'test', 5)
        self.data = self.data.reshape(-1, 3, 32, 32)
        save_cache_npz(self.data, self.label, url, self.train)

    def _load_data(self, filename, data_type='train', idx=None):
        """
        从文件中加载数据。

        Args:
            filename (str): 文件名。
            idx (int): 数据集索引。
            data_type (str, optional): 数据集类型。默认为 'train'。

        Returns:
            np.ndarray: 加载的数据。
        """
        assert data_type in ['train', 'test']
        with tarfile.open(filename, 'r:gz') as file:
            for item in file.getmembers():
                if ('data_batch_{}'.format(idx) in item.name and data_type == 'train') or (
                        'test_batch' in item.name and data_type == 'test'):
                    data_dict = pickle.load(file.extractfile(item), encoding='bytes')
                    data = data_dict[b'data']
                    return data

    def _load_label(self, filename, data_type='train', idx=None):
        """
        从文件中加载标签。

        Args:
            filename (str): 文件名。
            idx (int): 数据集索引。
            data_type (str, optional): 数据集类型。默认为 'train'。

        Returns:
            np.ndarray: 加载的标签。
        """
        assert data_type in ['train', 'test']
        with tarfile.open(filename, 'r:gz') as file:
            for item in file.getmembers():
                if ('data_batch_{}'.format(idx) in item.name and data_type == 'train') or (
                        'test_batch' in item.name and data_type == 'test'):
                    data_dict = pickle.load(file.extractfile(item), encoding='bytes')
                    return np.array(data_dict[b'labels'])

    def show(self, row=10, col=10):
        """
        显示 CIFAR-10 图像。

        Args:
            row (int, optional): 行数。默认为 10。
            col (int, optional): 列数。默认为 10。
        """
        H, W = 32, 32
        img = np.zeros((H * row, W * col, 3))
        for r in range(row):
            for c in range(col):
                img[r * H:(r + 1) * H, c * W:(c + 1) * W] = self.data[np.random.randint(0, len(self.data) - 1)].reshape(
                    3, H, W).transpose(1, 2, 0) / 255
        plt.imshow(img, interpolation='nearest')
        plt.axis('off')
        plt.show()

    @staticmethod
    def labels():
        """
        获取 CIFAR-10 的标签。

        Returns:
            dict: 标签字典。
        """
        return {0: 'ariplane', 1: 'automobile', 2: 'bird', 3: 'cat', 4: 'deer', 5: 'dog', 6: 'frog', 7: 'horse',
                8: 'ship', 9: 'truck'}
